from python_a2a.server import A2AServer
from python_a2a.models import Message, TextContent, MessageRole
from python_a2a import run_server
import re
import requests
import json

class YFinanceAgent(A2AServer):
    def __init__(self):
        super().__init__(url="http://localhost:5004")
        self.mcp_server_url = "http://localhost:5002"
    
    def handle_message(self, message):
        """同步处理消息"""
        if hasattr(message.content, 'text'):
            text = message.content.text
            
            # 提取股票代码
            ticker = self._extract_ticker(text)
            
            if not ticker:
                return Message(
                    content=TextContent(text="请提供股票代码来查询价格。"),
                    role=MessageRole.AGENT,
                    parent_message_id=getattr(message, 'message_id', None),
                    conversation_id=getattr(message, 'conversation_id', None)
                )
            
            # 使用MCP工具获取股票价格
            price_info = self._call_mcp_tool(ticker)
            import json
            price_data = json.loads(price_info['content'][0]['text'])
            price = price_data['price']
            currency = price_data.get("currency", "USD")
                
            return Message(
                content=TextContent(
                    text=f"{ticker} 当前交易价格为 {price:.2f} {currency}。"
                ),
                role=MessageRole.AGENT,
                parent_message_id=getattr(message, 'message_id', None),
                conversation_id=getattr(message, 'conversation_id', None)
            )
        
        return Message(
            content=TextContent(text="我可以提供股票代码的价格信息。"),
            role=MessageRole.AGENT,
            parent_message_id=getattr(message, 'message_id', None),
            conversation_id=getattr(message, 'conversation_id', None)
        )
    
    def _call_mcp_tool(self, ticker: str) -> dict:
        """调用MCP服务器获取股票价格"""
        payload = {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "tools/call",
            "params": {
                "name": "get_stock_price",
                "arguments": {
                    "ticker": ticker
                }
            }
        }
        
        try:
            response = requests.post(
                self.mcp_server_url,
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=10
            )
            
            if response.status_code == 200:
                result = response.json()
                
                if "result" in result:
                    price_info = result["result"]
                    return price_info
                elif "error" in result:
                    error_msg = f"MCP服务器错误: {result['error']}"
                    raise Exception(error_msg)
            
            raise Exception(f"HTTP错误: {response.status_code}")
            
        except Exception as e:
            raise e
    
    def _extract_ticker(self, text: str) -> str:
        ticker_match = re.search(r'\b([A-Z]{1,5})\b', text)
        if ticker_match:
            return ticker_match.group(1)
        
        patterns = [
            r"([A-Z]{1,5})的当前价格",
            r"([A-Z]{1,5})的价格",
            r"price of ([A-Z]{1,5})",
            r"price for ([A-Z]{1,5})"
        ]
        
        for pattern in patterns:
            match = re.search(pattern, text, re.IGNORECASE)
            if match:
                return match.group(1)
        
        return ""

if __name__ == "__main__":
    run_server(YFinanceAgent(), port=5004)